import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
from model import *
class CNNFramework:
	"""
	封装模型训练、评估与绘图功能。模型和数据分离，易于测试与复用。
	"""
	def __init__(self, model=None, input_shape=None, num_classes=None, pretrained_model_path=None):
		# 如果未提供现有模型，则重新构建模型
		if pretrained_model_path is None:
			self.model = model
			self.input_shape = input_shape
			self.num_classes = num_classes
		else:
			# 检查是否配置了预训练模型路径
			if pretrained_model_path and tf.io.gfile.exists(pretrained_model_path):
				print(f"加载预训练模型: {pretrained_model_path}")
				self.model = tf.keras.models.load_model(pretrained_model_path)
				# 添加权重检查和裁剪
				for layer in self.model.layers:
					weights = layer.get_weights()
					if weights:
						clipped_weights = []
						for w in weights:
							# 检查并打印权重范围
							print(f"Layer {layer.name}: max={np.max(w)}, min={np.min(w)}")
							# 裁剪到合理范围
							w = np.clip(w, -10, 10)
							clipped_weights.append(w)
						layer.set_weights(clipped_weights)
				
				self.model.summary()
				# 从预训练模型中提取输入形状和类别数
				self.input_shape = self.model.input_shape[1:]  # 去除batch维度
				self.num_classes = self.model.output_shape[-1]
			else:
				# 如果没有预训练模型，则创建新模型
				print("未找到预训练模型，创建新模型...")
				self.model = create_generic_model(input_shape, num_classes)
				self.model.summary()
				self.input_shape = input_shape
				self.num_classes = num_classes
		self.history = None

	def train_model(self, x_train,y_train, val_ds=None, epochs=10, callbacks=None):
		if callbacks is None:
			callbacks = []
		self.history = self.model.fit(x_train,
								y_train, 
								epochs=epochs,  
								verbose='auto',
								validation_data=val_ds, 
								callbacks=callbacks)
		print("训练完成。")

	def evaluate_model(self, test_ds):
		results = self.model.evaluate(test_ds)
		print(f"评估结果: {results}")
		return results
# 使用模型进行预测的示例
	def predict_new_data(model, new_data):
		"""
		使用训练好的模型进行预测

		参数:
		model: 训练好的模型
		new_data: 新的输入数据，形状为 (样本数, 时间步, 通道数)

		返回:
		predicted_classes: 预测的类别索引
		prediction_probabilities: 预测的概率值
		"""


		# 进行预测
		predictions = model.predict(new_data)

		# 获取预测结果
		predicted_classes = np.argmax(predictions, axis=1)
		prediction_probabilities = np.max(predictions, axis=1)

		return predicted_classes, prediction_probabilities, predictions
	def plot_training_history(self, metrics=['accuracy', 'loss'], plot_type='all_in_one'):
		"""
		绘制训练历史记录
		
		参数:
		metrics: 要绘制的指标列表，默认为['accuracy', 'loss']
		plot_type: 绘图类型，可选值为 'separate', 'all_in_one', 'split'，默认为'split'
		"""
		if self.history is None:
			print("没有可用的训练历史记录。")
			return
			
		if plot_type == 'separate':
			# 为每个指标分别绘制训练和验证曲线
			plt.figure(figsize=(10, 5 * max(1, len(metrics) // 2)))
			for i, metric in enumerate(metrics, 1):
				plt.subplot(len(metrics) // 2 + len(metrics) % 2, 2, i)
				if metric in self.history.history:
					plt.plot(self.history.history[metric], label=f'train {metric}')
				if f'val_{metric}' in self.history.history:
					plt.plot(self.history.history[f'val_{metric}'], label=f'val {metric}')
				plt.title(f'{metric} single change')
				plt.xlabel('training_times')
				plt.ylabel(metric)
				plt.legend()
			plt.tight_layout()
			plt.show()
		elif plot_type == 'all_in_one':
			# 将所有指标绘制在同一张图上
			plt.figure(figsize=(12, 6))
			
			# 为每个指标绘制训练和验证曲线
			for metric in metrics:
				# 绘制训练指标
				if metric in self.history.history:
					plt.plot(self.history.history[metric], label=f'train {metric}')
				
				# 绘制验证指标
				if f'val_{metric}' in self.history.history:
					plt.plot(self.history.history[f'val_{metric}'], label=f'val {metric}', linestyle='--')
			
			plt.title('Training Metrics')
			plt.xlabel('Epochs')
			plt.ylabel('Value')
			plt.legend()
			plt.grid(True)  # 添加网格线，使图表更易读
			plt.tight_layout()
			plt.show()
		elif plot_type == 'split':
			# 创建两个子图，一个用于训练数据，一个用于验证数据
			fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
			
			# 绘制训练指标
			for metric in metrics:
				if metric in self.history.history:
					ax1.plot(self.history.history[metric], label=metric)
			
			ax1.set_title('Training Metrics')
			ax1.set_xlabel('Epochs')
			ax1.set_ylabel('Value')
			ax1.legend()
			ax1.grid(True)
			
			# 绘制验证指标
			for metric in metrics:
				val_metric = f'val_{metric}'
				if val_metric in self.history.history:
					ax2.plot(self.history.history[val_metric], label=metric)
			
			ax2.set_title('Validation Metrics')
			ax2.set_xlabel('Epochs')
			ax2.set_ylabel('Value')
			ax2.legend()
			ax2.grid(True)
			
			plt.tight_layout()
			plt.show()
		else:
			print("无效的plot_type参数。请选择 'separate', 'all_in_one' 或 'split'。")